import sys

from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtCore import pyqtSignal
from PyQt5.QtGui import QMovie
from cnn_architecture import *
import inspect
import ctypes
import multiprocessing

import pickle


from keras.utils.generic_utils import Progbar


class UiMainWindow(QtWidgets.QMainWindow):
    show_mainwin_signal = pyqtSignal()
    show_main_input_signal = pyqtSignal()

    def __init__(self):
        super(UiMainWindow, self).__init__()
        self.setupUi(self)
        self.retranslateUi(self)

    def __getstate__(self):
        return self.__dict__

    def add_text(self, strr):
        self.textResult.append(strr+'\n')


    def setupUi(self, MainWindow):
        MainWindow.setObjectName("MainWindow")
        MainWindow.setEnabled(True)
        MainWindow.resize(518, 600)
        self.centralwidget = QtWidgets.QWidget(MainWindow)
        self.centralwidget.setEnabled(True)
        self.centralwidget.setObjectName("centralwidget")
        self.label = QtWidgets.QLabel(self.centralwidget)
        self.label.setGeometry(QtCore.QRect(0, 0, 611, 681))
        self.label.setText("")
        self.label.setPixmap(QtGui.QPixmap("333.png"))
        self.label.setObjectName("label")

        self.Button_create = QtWidgets.QPushButton(MainWindow)
        self.Button_create.setGeometry(QtCore.QRect(75, 85, 150, 61))
        self.Button_create.setText("开始训练")
        # self.Button_create.setFlat(True)
        self.Button_create.setObjectName("Button_create")
        self.Button_create.clicked.connect(self.train)

        self.Button_end = QtWidgets.QPushButton(MainWindow)
        self.Button_end.setGeometry(QtCore.QRect(300, 85, 150, 61))
        self.Button_end.setText("结束训练")
        # self.Button_end.setFlat(True)
        self.Button_end.setObjectName("Button_end")
        self.Button_end.setEnabled(False)
        self.Button_end.clicked.connect(self.endTrain)


        self.textResult = QtWidgets.QTextEdit(self.centralwidget)
        self.textResult.setGeometry(QtCore.QRect(0, 178, 521, 363))
        self.textResult.setObjectName("textResult")
        self.textResult.setReadOnly(True)

        self.label_gif = QtWidgets.QLabel(self.centralwidget)
        self.label_gif.setGeometry(QtCore.QRect(-70, 0, 671, 141))
        self.label_gif.setText("")
        self.label_gif.setObjectName("label_gif")
        MainWindow.setCentralWidget(self.centralwidget)
        self.menuBar = QtWidgets.QMenuBar(MainWindow)
        self.menuBar.setGeometry(QtCore.QRect(0, 0, 518, 22))
        self.menuBar.setObjectName("menuBar")
        self.menuCreate_Train = QtWidgets.QMenu(self.menuBar)
        self.menuCreate_Train.setObjectName("menuCreate_Train")
        self.menuV = QtWidgets.QMenu(self.menuBar)
        self.menuV.setObjectName("menuV")
        self.menuAccuracy = QtWidgets.QMenu(self.menuBar)
        self.menuAccuracy.setObjectName("menuAccuracy")
        self.menuResult = QtWidgets.QMenu(self.menuBar)
        self.menuResult.setObjectName("menuResult")
        self.menuAbout = QtWidgets.QMenu(self.menuBar)
        self.menuAbout.setObjectName("menuAbout")
        MainWindow.setMenuBar(self.menuBar)
        self.statusbar = QtWidgets.QStatusBar(MainWindow)
        self.statusbar.setObjectName("statusbar")
        MainWindow.setStatusBar(self.statusbar)
        self.toolBar = QtWidgets.QToolBar(MainWindow)
        self.toolBar.setObjectName("toolBar")
        MainWindow.addToolBar(QtCore.Qt.TopToolBarArea, self.toolBar)
        self.actionIn_data = QtWidgets.QAction(MainWindow)
        self.actionIn_data.setCheckable(False)
        self.actionIn_data.setObjectName("actionIn_data")
        self.actionIn_data.triggered.connect(self.to_input)#跳转到输入参数界面

        self.actionDew_P = QtWidgets.QAction(MainWindow)
        self.actionDew_P.setCheckable(False)
        self.actionDew_P.setObjectName("actionDew_P")
        self.actionDew_P.triggered.connect(self.to_dew)#跳转到泄漏点界面

        self.actionAccuracy = QtWidgets.QAction(MainWindow)
        self.actionAccuracy.setCheckable(False)
        self.actionAccuracy.setObjectName("actionAccuracy")
        self.actionAccuracy.triggered.connect(self.to_accuracy)

        self.actionExit = QtWidgets.QAction(MainWindow)
        self.actionExit.setCheckable(False)
        self.actionExit.setObjectName("actionExit")
        self.actionExit.triggered.connect(self.exit)

        self.actionVisualization = QtWidgets.QAction(MainWindow)
        self.actionVisualization.setObjectName("actionVisualization")
        self.actionVisualization.triggered.connect(self.to_visual)#跳转到可视化展示

        self.menuCreate_Train.addAction(self.actionIn_data)
        self.menuV.addAction(self.actionVisualization)
        self.menuAccuracy.addAction(self.actionAccuracy)
        self.menuResult.addAction(self.actionDew_P)
        self.menuAbout.addAction(self.actionExit)
        self.menuBar.addAction(self.menuCreate_Train.menuAction())
        self.menuBar.addAction(self.menuV.menuAction())
        self.menuBar.addAction(self.menuAccuracy.menuAction())
        self.menuBar.addAction(self.menuResult.menuAction())
        self.menuBar.addAction(self.menuAbout.menuAction())
        self.toolBar.addAction(self.actionIn_data)
        self.toolBar.addSeparator()
        self.toolBar.addAction(self.actionVisualization)
        self.toolBar.addSeparator()
        self.toolBar.addAction(self.actionDew_P)
        self.toolBar.addSeparator()
        self.toolBar.addAction(self.actionAccuracy)
        self.toolBar.addSeparator()
        self.toolBar.addAction(self.actionExit)

        self.retranslateUi(MainWindow)
        QtCore.QMetaObject.connectSlotsByName(MainWindow)

    def retranslateUi(self, MainWindow):
        _translate = QtCore.QCoreApplication.translate
        MainWindow.setWindowTitle(_translate("MainWindow", "MainWindow"))
        self.menuCreate_Train.setTitle(_translate("MainWindow", "Create Training"))
        self.menuV.setTitle(_translate("MainWindow", "Visualizion"))
        self.menuAccuracy.setTitle(_translate("MainWindow", "Accuracy"))
        self.menuResult.setTitle(_translate("MainWindow", "Result"))
        self.menuAbout.setTitle(_translate("MainWindow", "About"))
        self.toolBar.setWindowTitle(_translate("MainWindow", "toolBar"))
        self.actionIn_data.setText(_translate("MainWindow", "Input data"))
        self.actionDew_P.setText(_translate("MainWindow", "Dew Point"))
        self.actionAccuracy.setText(_translate("MainWindow", "Contrast"))
        self.actionExit.setText(_translate("MainWindow", "Exit"))
        self.actionVisualization.setText(_translate("MainWindow", "Visualization"))

    def train(self):
        # movie = QMovie("2.gif")
        # mainwin.label_gif.setMovie(movie)
        # movie.start()
        self.Button_create.setEnabled(False)
        self.Button_end.setEnabled(True)
        file = open("out.txt", 'w').close()  # 清空文件
        #cnn_thread.start()
        self.cnn_pro = multiprocessing.Process(target=cnn, args=(
        model_name, testKey, nb_attacks, nb_traces_attacks, learning_rate, batch_size, nb_epochs, real_key,
        input_size,))
        self.cnn_pro.start()
        P_info.start()

    def endTrain(self):
        self.Button_create.setEnabled(True)
        self.Button_end.setEnabled(False)
        self.cnn_pro.terminate()
        self.cnn_pro.join()

    def to_input(self):
     #   self.hide()
        inputwin.show()
        self.show_main_input_signal.emit()

    def to_visual(self):
        visualwin.show()

    def exit(self):
        self.close()

    def to_accuracy(self):
        accuwin.show()

    def to_dew(self):
        dewwin.show()


class Ui_first(QtWidgets.QWidget):

    show_mainwin_signal=pyqtSignal()

    def __init__(self):
        super(Ui_first,self).__init__()
        self.setupUi(self)
        self.retranslateUi(self)

    def setupUi(self, Form):
        Form.setObjectName("Form")
        Form.setEnabled(True)
        Form.resize(589, 759)
        self.label = QtWidgets.QLabel(Form)
        self.label.setGeometry(QtCore.QRect(0, -80, 591, 911))
        self.label.setText("")
        self.label.setPixmap(QtGui.QPixmap("first.png"))
        self.label.setObjectName("label")
        self.Button_in = QtWidgets.QPushButton(Form)
        self.Button_in.setGeometry(QtCore.QRect(160, 430, 271, 81))
        self.Button_in.setText("")
        self.Button_in.setDefault(False)
        self.Button_in.setFlat(True)
        self.Button_in.setObjectName("Button_in")
        self.Button_in.clicked.connect(self.to_main)#跳转到主窗口

        self.retranslateUi(Form)
        QtCore.QMetaObject.connectSlotsByName(Form)


    def retranslateUi(self, Form):
        _translate = QtCore.QCoreApplication.translate
        Form.setWindowTitle(_translate("Form", "侧信道分析平台"))

    def to_main(self):
        self.close()
        self.show_mainwin_signal.emit()


class Ui_input(QtWidgets.QWidget):

    show_main_input_signal = pyqtSignal()

    def __init__(self):
        super(Ui_input, self).__init__()
        self.setupUi(self)
        self.retranslateUi(self)


    def setupUi(self, widget):
        widget.setObjectName("widget")
        widget.setEnabled(True)
        widget.resize(567, 751)
        self.label = QtWidgets.QLabel(widget)
        self.label.setGeometry(QtCore.QRect(0, -5, 571, 761))
        self.label.setText("")
        self.label.setPixmap(QtGui.QPixmap("two2.png"))
        self.label.setObjectName("label")

        self.Button_data = QtWidgets.QPushButton(widget)
        self.Button_data.setGeometry(QtCore.QRect(161, 602, 270, 60))
        self.Button_data.setText("")
        self.Button_data.setFlat(True)
        self.Button_data.setObjectName("Button_data")
        self.line_name = QtWidgets.QLineEdit(widget)
        self.line_name.setGeometry(QtCore.QRect(320, 70, 221, 41))
        self.line_name.setObjectName("line_name")
        self.line_name.setText('AES_HD')
        self.line_testkey = QtWidgets.QLineEdit(widget)
        self.line_testkey.setGeometry(QtCore.QRect(320, 130, 221, 41))
        self.line_testkey.setObjectName("line_testkey")
        self.line_testkey.setText('0')
        self.line_attacknum = QtWidgets.QLineEdit(widget)
        self.line_attacknum.setGeometry(QtCore.QRect(320, 190, 221, 41))
        self.line_attacknum.setObjectName("line_attacknum")
        self.line_attacknum.setText('100')
        self.line_attacktrace = QtWidgets.QLineEdit(widget)
        self.line_attacktrace.setGeometry(QtCore.QRect(320, 250, 221, 41))
        self.line_attacktrace.setObjectName("line_attacktrace")
        self.line_attacktrace.setText('1500')
        self.line_lr = QtWidgets.QLineEdit(widget)
        self.line_lr.setGeometry(QtCore.QRect(320, 310, 221, 41))
        self.line_lr.setObjectName("line_lr")
        self.line_lr.setText('0.001')
        self.line_epochs = QtWidgets.QLineEdit(widget)
        self.line_epochs.setGeometry(QtCore.QRect(320, 370, 221, 41))
        self.line_epochs.setObjectName("line_epochs")
        self.line_epochs.setText('10')
        self.line_size = QtWidgets.QLineEdit(widget)
        self.line_size.setGeometry(QtCore.QRect(320, 430, 221, 41))
        self.line_size.setObjectName("line_size")
        self.line_size.setText('256')
        self.line_node = QtWidgets.QLineEdit(widget)
        self.line_node.setGeometry(QtCore.QRect(320, 490, 221, 41))
        self.line_node.setObjectName("line_node")
        self.line_node.setText('1250')

        self.retranslateUi(widget)
        QtCore.QMetaObject.connectSlotsByName(widget)

        self.Button_data.clicked.connect(self.load)

    def retranslateUi(self, widget):
        _translate = QtCore.QCoreApplication.translate
        widget.setWindowTitle(_translate("widget", "参数输入"))

    def load(self):
        global model_name,testKey,nb_epochs,nb_epochs,nb_attacks,nb_traces_attacks,learning_rate,batch_size,input_size
        model_name=self.line_name.text()
        testKey=int(self.line_testkey.text())
        nb_attacks=int(self.line_attacknum.text())
        nb_traces_attacks=int(self.line_attacktrace.text())
        learning_rate=float(self.line_lr.text())
        nb_epochs=int(self.line_epochs.text())
        batch_size=int(self.line_size.text())
        input_size=int(self.line_node.text())
        self.close()

class Ui_att(QtWidgets.QWidget):
    def __init__(self):
        super(Ui_att,self).__init__()
        self.setupUi(self)
        self.retranslateUi(self)
    def setupUi(self,Form):
        Form.setObjectName("att_Form")
        Form.resize(1250,630)
        self.label_att = QtWidgets.QLabel(Form)
        self.label_att.setGeometry(QtCore.QRect(390, 120, 470, 350))
        self.label_att.setText("")

        self.label_att.setObjectName("label_att")

    def retranslateUi(self, Form):
        _translate = QtCore.QCoreApplication.translate
        Form.setWindowTitle(_translate("Form", "att_Form"))


class Ui_visual(QtWidgets.QWidget):
    def __init__(self):
        super(Ui_visual,self).__init__()
        self.setupUi(self)
        self.retranslateUi(self)
        self.ui_att = Ui_att()
        self.ui_att.hide()

    def setupUi(self, Form):
        Form.setObjectName("Form")
        Form.resize(1920, 1000)
        #self.label = QtWidgets.QLabel(Form)
        #self.label.setGeometry(QtCore.QRect(0, 20, 601, 781))
        #self.label.setText("")
        #self.label.setPixmap(QtGui.QPixmap("4.png"))
        #self.label.setObjectName("label")

        self.label_f = QtWidgets.QLabel(Form)
        self.label_f.setGeometry(QtCore.QRect(0, 20, 640, 480))
        self.label_f.setText("")

        self.label_f.setObjectName("label_f")

        self.label_g = QtWidgets.QLabel(Form)
        self.label_g.setGeometry(QtCore.QRect(1280, 20, 640, 480))
        self.label_g.setText("")

        self.label_g.setObjectName("label_g")

        # button
        self.Button_att = QtWidgets.QPushButton(Form)
        self.Button_att.setGeometry(QtCore.QRect(750, 240, 341, 61))
        self.Button_att.setText("攻击Rank曲线")
        # self.Button_att.setFlat(True)  # 去边框
        self.Button_att.setObjectName("Button_att")

        self.label_w1 = QtWidgets.QLabel(Form)
        self.label_w1.setGeometry(QtCore.QRect(0, 520, 640, 480))
        self.label_w1.setText("")

        self.label_w1.setObjectName("label_w1")

        self.label_w2 = QtWidgets.QLabel(Form)
        self.label_w2.setGeometry(QtCore.QRect(640, 520, 640, 480))
        self.label_w2.setText("")

        self.label_w2.setObjectName("label_w2")

        self.label_w3 = QtWidgets.QLabel(Form)
        self.label_w3.setGeometry(QtCore.QRect(1280, 520, 640, 480))
        self.label_w3.setText("")

        self.label_w3.setObjectName("label_w3")

        self.retranslateUi(Form)
        QtCore.QMetaObject.connectSlotsByName(Form)

    def retranslateUi(self, Form):
        _translate = QtCore.QCoreApplication.translate
        Form.setWindowTitle(_translate("Form", "Form"))
        self.Button_att.clicked.connect(self.show_att)

    def show_att(self):
        self.ui_att.show()

    def show_graph(self):
        self.label_f.setPixmap(QtGui.QPixmap("feather_visual_" + str(nb_epochs) + '.png'))
        self.label_g.setPixmap(QtGui.QPixmap("gradient_visual_" + str(nb_epochs) + '.png'))
        self.label_w1.setPixmap(QtGui.QPixmap("weight_visual_" + str(nb_epochs) + '_0' + '.png'))
        self.label_w2.setPixmap(QtGui.QPixmap("weight_visual_" + str(nb_epochs) + '_1' + '.png'))
        self.label_w3.setPixmap(QtGui.QPixmap("weight_visual_" + str(nb_epochs) + '_ALL' + '.png'))
        self.ui_att.label_att.setPixmap(QtGui.QPixmap('fig/rank' + model_name + '_' + str(nb_traces_attacks) + 'trs_'
                                                      + str(nb_attacks) + 'att.svg'))

        # 获取可视化图片
        #特征可视化（"feather_visual_" + str(Thread_cnn.nb_epochs) + '.png'）
    def get_picture(self):
        print(1)
        #读取图片
        #显示图片


        #权重可视化（"weight_visual_" + str(Thread_cnn.nb_epochs) + '_ALL' + '.png'）
    def get_picture(self):
        print(1)
        # 读取图片
        # 显示图片


        #梯度可视化（"gradient_visual_" + str(Thread_cnn.nb_epochs) + '.png'）
    def get_picture(self):
        print(1)
        # 读取图片
        # 显示图片




class Ui_Accuracy(QtWidgets.QWidget):
    def __init__(self):
        super(Ui_Accuracy,self).__init__()
        self.setupUi(self)
        self.retranslateUi(self)
    def setupUi(self, Form):
        Form.setObjectName("Form")
        Form.resize(589, 762)
        self.label = QtWidgets.QLabel(Form)
        self.label.setGeometry(QtCore.QRect(0, -5, 621, 771))
        self.label.setText("")
        self.label.setPixmap(QtGui.QPixmap("准确率.png"))
        self.label.setObjectName("label")
        self.line_gra = QtWidgets.QLineEdit(Form)
        self.line_gra.setGeometry(QtCore.QRect(290, 210, 211, 41))
        self.line_gra.setObjectName("line_gra")
        self.line_gra.setEnabled(False)
        self.line_nongra = QtWidgets.QLineEdit(Form)
        self.line_nongra.setGeometry(QtCore.QRect(290, 420, 211, 41))
        self.line_nongra.setObjectName("line_nongra")
        self.line_nongra.setEnabled(False)

        self.retranslateUi(Form)
        QtCore.QMetaObject.connectSlotsByName(Form)

    def retranslateUi(self, Form):
        _translate = QtCore.QCoreApplication.translate
        Form.setWindowTitle(_translate("Form", "准确率比较"))

    def setText(self):
        file1 = open('fig/sum.txt', 'r')
        file2 = open('fig/sum_New.txt', 'r')
        str1=''
        str2=''
        while 1:
            char = file1.read(1)
            if char == ',':
                break
            else:
                str1 = str1 + char
            time.sleep(0.01)
        while 1:
            char = file2.read(1)
            if char == ',':
                break
            else:
                str2 = str2 + char
            time.sleep(0.01)
        str1 = str1 + '/' + str(nb_traces_attacks)
        str2 = str2 + '/' + str(nb_traces_attacks)
        self.line_gra.setText(str2)
        self.line_nongra.setText(str1)

class Ui_Dew(QtWidgets.QWidget):
    def __init__(self):
        super(Ui_Dew, self).__init__()
        self.setupUi(self)
        self.retranslateUi(self)

    def setupUi(self, Form):
        Form.setObjectName("Form")
        Form.resize(338, 273)
        self.label = QtWidgets.QLabel(Form)
        self.label.setGeometry(QtCore.QRect(0, -10, 341, 291))
        self.label.setText("")
        self.label.setPixmap(QtGui.QPixmap("泄漏点.png"))
        self.label.setObjectName("label")
        self.line_low = QtWidgets.QLineEdit(Form)
        self.line_low.setGeometry(QtCore.QRect(60, 150, 41, 41))
        self.line_low.setObjectName("text_low")
        self.line_low.setEnabled(False)
        self.line_high = QtWidgets.QLineEdit(Form)
        self.line_high.setGeometry(QtCore.QRect(220, 150, 41, 41))
        self.line_high.setObjectName("text_high")
        self.line_high.setEnabled(False)

        self.retranslateUi(Form)
        QtCore.QMetaObject.connectSlotsByName(Form)

    def retranslateUi(self, Form):
        _translate = QtCore.QCoreApplication.translate
        Form.setWindowTitle(_translate("Form", "Dew pint"))

    def setText(self):
        file=open('dew.txt','r')
        strr = file.readline()
        str_low = strr[(strr.find('low_i:') + 6):(strr.find('high_i:') - 1)]
        str_high = strr[(strr.find('high_i:') + 7):]
        self.line_low.setText(str_low)
        self.line_high.setText(str_high)

def show_main():

    win.hide()
    mainwin.show()


def show_input():
    inputwin.show()


def _async_raise(tid, exctype):
    """raises the exception, performs cleanup if needed"""
    tid = ctypes.c_long(tid)
    if not inspect.isclass(exctype):
        exctype = type(exctype)
    res = ctypes.pythonapi.PyThreadState_SetAsyncExc(tid, ctypes.py_object(exctype))
    if res == 0:
        raise ValueError("invalid thread id")
    elif res != 1:
        # """if it returns a number greater than one, you're in trouble,
        # and you should call it again with exc=NULL to revert the effect"""
        ctypes.pythonapi.PyThreadState_SetAsyncExc(tid, None)
        raise SystemError("PyThreadState_SetAsyncExc failed")
def stop_thread(thread):
    _async_raise(thread.ident, SystemExit)




class PrintInfo(threading.Thread):
    def __init__(self, obj1):
        threading.Thread.__init__(self)
        self.obj1 = obj1

    def run(self):
        file = open('out.txt', 'r')
        while 1:
            str = file.readline()
            if not str:
                time.sleep(2)  # 2s
                continue
            else:
                str=str.strip('\n')
                self.obj1.add_text(str)
                if str == '权重可视化已输出' or str == '特征可视化已输出' or str == '梯度可视化已输出' or str == '攻击Rank曲线已输出':
                    visualwin.show_graph()
                if str.find('low_i:') != -1 and str.find('high_i:') != -1:
                    #print()     # 得到泄露点
                    dewwin.setText()
                if str == '训练结束，请查看结果！！':
                    #print('1')       # 计算梯度/非梯度反馈准确率
                    accuwin.setText()

                time.sleep(0.01)


if __name__=="__main__":


    nb_epochs = 10  # 过10次所有训练样本  # 论文20
    batch_size = 256  # 一次训练放256个样本
    input_size = 1250  # X轴   一条功耗曲线有1250个样本点（时间点）
    learning_rate = 1e-3  # 0.001
    nb_traces_attacks = 1500  # 攻击一次要用1500条曲线
    nb_attacks = 100  # 100次攻击
    real_key = [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]
    testKey = 0  # 攻击字节
    model_name = 'AES_HD'



    app=QtWidgets.QApplication(sys.argv)
    win=Ui_first()

    mainwin=UiMainWindow()

    inputwin=Ui_input()

    visualwin=Ui_visual()

    accuwin=Ui_Accuracy()

    dewwin=Ui_Dew()

    P_info = PrintInfo(mainwin)
    #cnn_thread = Thread_cnn()
    #cnn_thread.tran(mainwin)



    win.show()
    win.show_mainwin_signal.connect(show_main)
    inputwin.show_main_input_signal.connect(show_input)
    sys.exit(app.exec_())
